//-------------------------------------------------------------
//
// HLDNS: this is the implementation of hldns hooks into tinydns.
//        There are three hooks, init, lookup and exit.
//        The primary purpose is to tell tinydns which record DB
//        to answer DNS requests from based on the requester IP and
//        the Anycast shortest route (implied by the fact that the
//        request made it to the local server). It also logs these
//        results for refining of the algorithm and regionalization
//        over time.
//
// Author: Jason Rexilius <jason@hostedlabs.com>
//         John Quigley <jquigley@jquigley.com>
//         Tristan Sloughter <kungfooguru@gmail.com>
//
//-------------------------------------------------------------

//-------------------------------------------------------------
// REQUIRED LIBRARIES
//-------------------------------------------------------------
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <time.h>
#include <arpa/inet.h>
#include <GeoIP.h>
#include <GeoIPCity.h>

//-------------------------------------------------------------
// VARIABLES AND DEFAULTS
//-------------------------------------------------------------

// ENV VARS SHOULD OVERWRITE THESE
char *DEFAULT_LOCAL_AOR = "CENTRAL.cdb";
char *DEFAULT_LOCAL_COUNTRY = "US";
char *DEFAULT_LOCAL_REGION = "IL";
char *DEFAULT_HOSTNAME = "ns0";
char *DEFAULT_LOG_FILE_PATH = "hldns.log";
FILE *log_output_stream = NULL;
char *DEFAULT_GEOIP_DB_PATH = "/usr/share/GeoIP/GeoIPCity.dat";

// ENV VAR CONTAINER
typedef struct _SystemConfig{
 // GeoIP
 GeoIP *geoip;
 int geoip_open;
 // local data
 char *local_aor;
 char *local_country;
 char *local_region;
 char *hostname;
 // paths
 char *geoip_db_path;
 char *log_file_path;
} SystemConfig;
SystemConfig sys;

//-------------------------------------------------------------
// AOR TABLE - sorted by traffic/population for US web sites
//-------------------------------------------------------------
const char AOR[][3][15] = {
"US","CA","WEST.cdb",
"US","NY","EAST.cdb",
"US","TX","CENTRAL.cdb",
"US","IL","CENTRAL.cdb",
"US","AA","CENTRAL.cdb",
"US","AL","CENTRAL.cdb",
"US","AR","CENTRAL.cdb",
"US","IA","CENTRAL.cdb",
"US","IN","CENTRAL.cdb",
"US","KS","CENTRAL.cdb",
"US","KY","CENTRAL.cdb",
"US","LA","CENTRAL.cdb",
"US","MI","CENTRAL.cdb",
"US","MN","CENTRAL.cdb",
"US","MO","CENTRAL.cdb",
"US","MS","CENTRAL.cdb",
"US","ND","CENTRAL.cdb",
"US","NE","CENTRAL.cdb",
"US","OH","CENTRAL.cdb",
"US","OK","CENTRAL.cdb",
"US","SD","CENTRAL.cdb",
"US","TN","CENTRAL.cdb",
"US","WI","CENTRAL.cdb",
"US","","CENTRAL.cdb",
"US","CT","EAST.cdb",
"US","DC","EAST.cdb",
"US","DE","EAST.cdb",
"US","FL","EAST.cdb",
"US","GA","EAST.cdb",
"US","MA","EAST.cdb",
"US","MD","EAST.cdb",
"US","ME","EAST.cdb",
"US","NC","EAST.cdb",
"US","NH","EAST.cdb",
"US","NJ","EAST.cdb",
"US","NY","EAST.cdb",
"US","PA","EAST.cdb",
"US","PR","EAST.cdb",
"US","RI","EAST.cdb",
"US","SC","EAST.cdb",
"US","VA","EAST.cdb",
"US","VI","EAST.cdb",
"US","VT","EAST.cdb",
"US","WV","EAST.cdb",
"US","AK","WEST.cdb",
"US","AZ","WEST.cdb",
"US","CO","WEST.cdb",
"US","HI","WEST.cdb",
"US","ID","WEST.cdb",
"US","MT","WEST.cdb",
"US","NM","WEST.cdb",
"US","NV","WEST.cdb",
"US","OR","WEST.cdb",
"US","UT","WEST.cdb",
"US","WA","WEST.cdb",
"US","WY","WEST.cdb",
"US","AE","EU.cdb",
"US","AP","APAC.cdb",
"US","AS","APAC.cdb",
"US","FM","APAC.cdb",
"US","GU","APAC.cdb",
"US","MH","APAC.cdb",
"US","MP","APAC.cdb",
"US","PW","APAC.cdb",
"CA","","CENTRAL.cdb",
"CA","MB","CENTRAL.cdb",
"CA","NU","CENTRAL.cdb",
"CA","ON","CENTRAL.cdb",
"CA","SK","CENTRAL.cdb",
"CA","NB","EAST.cdb",
"CA","NL","EAST.cdb",
"CA","NS","EAST.cdb",
"CA","PE","EAST.cdb",
"CA","QC","EAST.cdb",
"CA","AB","WEST.cdb",
"CA","BC","WEST.cdb",
"CA","NT","WEST.cdb",
"CA","YT","WEST.cdb",
"DE","","EU.cdb",
"FR","","EU.cdb",
"GB","","EU.cdb",
"UK","","EU.cdb",
"IT","","EU.cdb",
"ES","","EU.cdb",
"VI","","EAST.cdb",
"AD","","EU.cdb",
"AL","","EU.cdb",
"AM","","EU.cdb",
"AT","","EU.cdb",
"AX","","EU.cdb",
"AZ","","EU.cdb",
"BA","","EU.cdb",
"BD","","EU.cdb",
"BE","","EU.cdb",
"BG","","EU.cdb",
"BY","","EU.cdb",
"CH","","EU.cdb",
"CZ","","EU.cdb",
"DK","","EU.cdb",
"EE","","EU.cdb",
"EU","","EU.cdb",
"FI","","EU.cdb",
"FO","","EU.cdb",
"GE","","EU.cdb",
"GG","","EU.cdb",
"GI","","EU.cdb",
"GL","","EU.cdb",
"GR","","EU.cdb",
"HR","","EU.cdb",
"HU","","EU.cdb",
"IE","","EU.cdb",
"IM","","EU.cdb",
"IS","","EU.cdb",
"JE","","EU.cdb",
"KG","","EU.cdb",
"KZ","","EU.cdb",
"LI","","EU.cdb",
"LT","","EU.cdb",
"LU","","EU.cdb",
"LV","","EU.cdb",
"MC","","EU.cdb",
"MD","","EU.cdb",
"ME","","EU.cdb",
"MK","","EU.cdb",
"MT","","EU.cdb",
"NL","","EU.cdb",
"NO","","EU.cdb",
"PL","","EU.cdb",
"PT","","EU.cdb",
"RO","","EU.cdb",
"RS","","EU.cdb",
"RU","","EU.cdb",
"SE","","EU.cdb",
"SH","","EU.cdb",
"SI","","EU.cdb",
"AI","","EAST.cdb",
"BB","","EAST.cdb",
"BM","","EAST.cdb",
"BS","","EAST.cdb",
"DM","","EAST.cdb",
"DO","","EAST.cdb",
"GP","","EAST.cdb",
"HT","","EAST.cdb",
"JM","","EAST.cdb",
"KN","","EAST.cdb",
"KY","","EAST.cdb",
"LC","","EAST.cdb",
"MQ","","EAST.cdb",
"MS","","EAST.cdb",
"PM","","EAST.cdb",
"PR","","EAST.cdb",
"TC","","EAST.cdb",
"VC","","EAST.cdb",
"VG","","EAST.cdb",
"SJ","","EU.cdb",
"SK","","EU.cdb",
"SM","","EU.cdb",
"TJ","","EU.cdb",
"TM","","EU.cdb",
"TR","","EU.cdb",
"UA","","EU.cdb",
"UZ","","EU.cdb",
"VA","","EU.cdb",
"JP","","APAC.cdb",
"KR","","APAC.cdb",
"CN","","APAC.cdb",
"IN","","MEAFR.cdb",
"AP","","APAC.cdb",
"AQ","","APAC.cdb",
"AS","","APAC.cdb",
"AU","","APAC.cdb",
"BN","","APAC.cdb",
"BT","","APAC.cdb",
"CC","","APAC.cdb",
"CK","","APAC.cdb",
"CX","","APAC.cdb",
"FJ","","APAC.cdb",
"FM","","APAC.cdb",
"GU","","APAC.cdb",
"HK","","APAC.cdb",
"HM","","APAC.cdb",
"ID","","APAC.cdb",
"IO","","APAC.cdb",
"KH","","APAC.cdb",
"KI","","APAC.cdb",
"KP","","APAC.cdb",
"LA","","APAC.cdb",
"LK","","APAC.cdb",
"MH","","APAC.cdb",
"MM","","APAC.cdb",
"MN","","APAC.cdb",
"MO","","APAC.cdb",
"MP","","APAC.cdb",
"MY","","APAC.cdb",
"NC","","APAC.cdb",
"NF","","APAC.cdb",
"NP","","APAC.cdb",
"NR","","APAC.cdb",
"NU","","APAC.cdb",
"NZ","","APAC.cdb",
"PF","","APAC.cdb",
"PG","","APAC.cdb",
"PH","","APAC.cdb",
"PN","","APAC.cdb",
"PW","","APAC.cdb",
"SB","","APAC.cdb",
"SG","","APAC.cdb",
"TF","","APAC.cdb",
"TH","","APAC.cdb",
"TK","","APAC.cdb",
"TL","","APAC.cdb",
"TO","","APAC.cdb",
"TV","","APAC.cdb",
"TW","","APAC.cdb",
"UM","","APAC.cdb",
"VN","","APAC.cdb",
"VU","","APAC.cdb",
"WF","","APAC.cdb",
"WS","","APAC.cdb",
"AG","","LATIN.cdb",
"AN","","LATIN.cdb",
"AR","","LATIN.cdb",
"AW","","LATIN.cdb",
"BO","","LATIN.cdb",
"BR","","LATIN.cdb",
"BZ","","LATIN.cdb",
"CL","","LATIN.cdb",
"CO","","LATIN.cdb",
"CR","","LATIN.cdb",
"CU","","LATIN.cdb",
"EC","","LATIN.cdb",
"FK","","LATIN.cdb",
"GD","","LATIN.cdb",
"GF","","LATIN.cdb",
"GS","","LATIN.cdb",
"GT","","LATIN.cdb",
"GY","","LATIN.cdb",
"HN","","LATIN.cdb",
"MX","","LATIN.cdb",
"NI","","LATIN.cdb",
"PA","","LATIN.cdb",
"PE","","LATIN.cdb",
"PY","","LATIN.cdb",
"SR","","LATIN.cdb",
"SV","","LATIN.cdb",
"TT","","LATIN.cdb",
"UY","","LATIN.cdb",
"VE","","LATIN.cdb",
"AE","","MEAFR.cdb",
"AF","","MEAFR.cdb",
"AO","","MEAFR.cdb",
"BF","","MEAFR.cdb",
"BH","","MEAFR.cdb",
"BI","","MEAFR.cdb",
"BJ","","MEAFR.cdb",
"BV","","MEAFR.cdb",
"BW","","MEAFR.cdb",
"CD","","MEAFR.cdb",
"CF","","MEAFR.cdb",
"CG","","MEAFR.cdb",
"CI","","MEAFR.cdb",
"CM","","MEAFR.cdb",
"CV","","MEAFR.cdb",
"CY","","MEAFR.cdb",
"DJ","","MEAFR.cdb",
"DZ","","MEAFR.cdb",
"EG","","MEAFR.cdb",
"EH","","MEAFR.cdb",
"ER","","MEAFR.cdb",
"ET","","MEAFR.cdb",
"GA","","MEAFR.cdb",
"GH","","MEAFR.cdb",
"GM","","MEAFR.cdb",
"GN","","MEAFR.cdb",
"GQ","","MEAFR.cdb",
"GW","","MEAFR.cdb",
"IL","","MEAFR.cdb",
"IQ","","MEAFR.cdb",
"IR","","MEAFR.cdb",
"JO","","MEAFR.cdb",
"KE","","MEAFR.cdb",
"KM","","MEAFR.cdb",
"KW","","MEAFR.cdb",
"LB","","MEAFR.cdb",
"LR","","MEAFR.cdb",
"LS","","MEAFR.cdb",
"LY","","MEAFR.cdb",
"MA","","MEAFR.cdb",
"MG","","MEAFR.cdb",
"ML","","MEAFR.cdb",
"MR","","MEAFR.cdb",
"MU","","MEAFR.cdb",
"MV","","MEAFR.cdb",
"MW","","MEAFR.cdb",
"MZ","","MEAFR.cdb",
"NA","","MEAFR.cdb",
"NE","","MEAFR.cdb",
"NG","","MEAFR.cdb",
"OM","","MEAFR.cdb",
"PK","","MEAFR.cdb",
"PS","","MEAFR.cdb",
"QA","","MEAFR.cdb",
"RE","","MEAFR.cdb",
"RW","","MEAFR.cdb",
"SA","","MEAFR.cdb",
"SC","","MEAFR.cdb",
"SD","","MEAFR.cdb",
"SL","","MEAFR.cdb",
"SN","","MEAFR.cdb",
"SO","","MEAFR.cdb",
"ST","","MEAFR.cdb",
"SY","","MEAFR.cdb",
"SZ","","MEAFR.cdb",
"TD","","MEAFR.cdb",
"TG","","MEAFR.cdb",
"TN","","MEAFR.cdb",
"TZ","","MEAFR.cdb",
"UG","","MEAFR.cdb",
"YE","","MEAFR.cdb",
"YT","","MEAFR.cdb",
"ZA","","MEAFR.cdb",
"ZM","","MEAFR.cdb",
"ZW","","MEAFR.cdb"
};


//-------------------------------------------------------------
// INITIALIZATION HOOK - sets up environment, log and geoip
//-------------------------------------------------------------

void hooks_server_init(void)
{
  // OPEN LOG FILE
  //  if fail to open use stdout
  sys.log_file_path = getenv("LOG_FILE_PATH");
  if (sys.log_file_path == NULL){
   sys.log_file_path = DEFAULT_LOG_FILE_PATH;
   fprintf(stderr,"WARN: ENV VAR NOT SET, LOG_FILE_PATH SET TO DEFAULT %s\n",DEFAULT_LOG_FILE_PATH);
  }
  log_output_stream = fopen(sys.log_file_path, "a");
  if (log_output_stream == NULL){
   fprintf(stderr,"ERROR: FAILED TO OPEN log_output_stream USING %s\n",sys.log_file_path);
   log_output_stream = stdout;
  }
  
  // READ IN ENV VARS
  //  if no env vars, log error, continue with defaults
  sys.local_aor = getenv("LOCAL_AOR");
  if (sys.local_aor == NULL){
   sys.local_aor = DEFAULT_LOCAL_AOR;
   fprintf(stderr,"WARN: ENV VAR NOT SET, LOCAL_AOR USING %s\n",DEFAULT_LOCAL_AOR);
  }

  sys.local_country = getenv("LOCAL_COUNTRY");
  if (sys.local_country == NULL){
   sys.local_country = DEFAULT_LOCAL_COUNTRY;
   fprintf(stderr,"WARN: ENV VAR NOT SET, LOCAL_COUNTRY USING %s\n",DEFAULT_LOCAL_COUNTRY);
  }

  sys.local_region = getenv("LOCAL_REGION");
  if (sys.local_region == NULL){ 
   sys.local_region = DEFAULT_LOCAL_REGION;
   fprintf(stderr,"WARN: ENV VAR NOT SET, LOCAL_REGION USING %s\n",DEFAULT_LOCAL_REGION);
  }

  sys.hostname = getenv("HOSTNAME");
  if (sys.hostname == NULL){
   sys.hostname = DEFAULT_HOSTNAME;
   fprintf(stderr,"WARN: ENV VAR NOT SET, HOSTNAME USING %s\n",DEFAULT_HOSTNAME);
  }

  sys.geoip_db_path = getenv("GEOIP_DB_PATH");
  if (sys.geoip_db_path == NULL){
    sys.geoip_db_path = DEFAULT_GEOIP_DB_PATH;  
    fprintf(stderr,"WARN: ENV VAR NOT SET, GEOIP_DB_PATH  USING %s\n",DEFAULT_GEOIP_DB_PATH);
  }

  // INIT GEOIP
  sys.geoip = GeoIP_open(sys.geoip_db_path,GEOIP_MEMORY_CACHE);
  if (sys.geoip != NULL){
   sys.geoip_open=0;
  }else{
   fprintf (stderr,"ERROR: FAILED TO OPEN GEOIPDB: %s\n", sys.geoip_db_path);
   sys.geoip_open=1;
  }

  return;
}

//-------------------------------------------------------------
// MAIN HOOK - perform lookup and determine which table to use
//-------------------------------------------------------------

int hooks_tdlookup_open_cdb(char ip[4], char *cdb_file_path, int path_len)
{
 // INTERFACE AND INPUT VARS
 char reqdom[255]="reqdomain.com";

 // INIT GEOIP RECORD
 GeoIPRecord *geoip_record;

 // HOW MANY AOR RECORDS ARE THERE?
 int count = 0;
 int elems = 0;
 elems=sizeof AOR /sizeof AOR[0];
 if(elems < 1){fprintf(stderr,"ERROR: elems=sizeof AOR /sizeof AOR[0] YIELDS LESS THAN 1 ??\n");}

 // MAKE IP STRING FROM IP NUM
 char ip_str[16];
 int ip_len=sizeof(ip_str);
 int fresult=0;
 memset(ip_str,'\0', ip_len);
 fresult=ip2str(ip, ip_str, ip_len);

 // CALL GEOIP, FETCH COUNTRY_CODE, REGION_CODE
 char cc[15];
 memset(cc,'\0',sizeof(cc));
 char st[15];
 memset(st,'\0',sizeof(st));

 // DO GEOIP LOOKUP, EXPECT VALUES OR EMPTY STRING
 if(sys.geoip_open==0){
  geoip_record = GeoIP_record_by_addr(sys.geoip, (const char *)ip_str);
  if (geoip_record != NULL) {
   if (geoip_record->country_code != NULL){
    strncpy (cc, geoip_record->country_code, strlen(geoip_record->country_code));
    if (geoip_record->region != NULL){
     strncpy (st, geoip_record->region, strlen(geoip_record->region));
    }
   }
   GeoIPRecord_delete(geoip_record);
  }
 }

 // NULL OUT REGION FOR ALL BUT US AND CANADA
 if( strcmp(cc,"US")!=0 && strcmp(cc,"US")!=0 ){memset(st,'\0',sizeof(st));}

 // IF GEOIP LOOKUP FAILED TO FIND COUNTRY, USE LOCAL
 if(strcmp(cc,"")==0){
  strncpy(cdb_file_path,sys.local_aor,path_len);
  fresult=log_record(1,ip_str,reqdom,cdb_file_path);
  return 0;
 }

 // LOOKUP AOR FROM TABLE, IF CC AND ST MATCH, USE AOR
 while(count < elems){
  if( strcmp(cc,AOR[count][0]) == 0 && strcmp(st,AOR[count][1]) == 0 ){
   strncpy(cdb_file_path,AOR[count][2],path_len);
   fresult=log_record(0,ip_str,reqdom,cdb_file_path);
   return 0;
  }
  count++;
 }
 
 // DEFAULT TO LOCAL AOR
 strncpy(cdb_file_path,sys.local_aor,path_len);
 fresult=log_record(2,ip_str,reqdom,cdb_file_path);
 return 0;
}

//-------------------------------------------------------------
// EXIT HOOK - cleanup log and geoip
//-------------------------------------------------------------

void hooks_server_exit(void)
{
 // CLENUP GEOIP
 GeoIP_delete(sys.geoip);

 // CLOSE LOG FILE
 fclose(log_output_stream);
 return;
}

//-------------------------------------------------------------
// MISCELLANEOUS FUNCTIONS
//-------------------------------------------------------------

int ip2str(const char *ip, char *str, int str_len){
 struct in_addr addr;
 memcpy(&addr.s_addr, ip, sizeof(addr.s_addr));
 strncpy(str, inet_ntoa(addr), str_len);
 return 0;
}

int log_record(int retval,const char *ip_str, const char *reqdom, const char *aor){
 time_t time_now = time(NULL);
 struct tm *now  = localtime(&time_now);
 FILE *s = log_output_stream;
 if (s == NULL){s = stdout;}
 fprintf(s, "%i-%.2i-%.2i-%.2i-%.2i-%.2i,%s,%s,%s,%d,%s\n",
  now->tm_year + 1900,
  now->tm_mon + 1,
  now->tm_mday,
  now->tm_hour,
  now->tm_min,
  now->tm_sec,
  sys.hostname,
  ip_str,
  reqdom,
  retval,
  aor
  );
  fflush(s);
  return;
} 

